//=======================================================================================
// HMM.cc
// --------------------------------------------------------------------------------------
//
//---------------------------------------------------------------------------------------
// $Id: HMM.cc,v 1.1 2000/01/31 21:14:27 kdurbin Exp $
//---------------------------------------------------------------------------------------
// $Log: HMM.cc,v $
// Revision 1.1  2000/01/31 21:14:27  kdurbin
// Initial revision
//
//   
//========================================================================================
//.......1.........2.........3.........4.........5.........6.........7.........8.........9
//

#include "HMM.hh"							

// Something close to zero in log space.  log10(0) is undefined, so we can't
// ever have exactly zero in our calculations in log space. 
#define ROUGHLY_ZERO -1000000000

char
HMM::bugChar(int symbol)
{
	char c[4]={'A','C','G','T'};
	if (symbol > 3) {
	  cout<<"Invalid symbol number."<<endl;
	  return('X');
	}
	return(c[symbol]);
}

//---------------------------------------------------------------------------
// * BaseIdx
//---------------------------------------------------------------------------
// Converts ASCII letters into an index. 
// 
int
HMM::BaseIdx(char c)
{
	switch(c){
		case 'A':
		case 'a':
			return(0);
			break;
		case 'C':
		case 'c':
			return(1);
			break;
		case 'G':
		case 'g':
			return(2);
			break;
		case 'T':
		case 't':
			return(3);
			break;
		case 'N':
		case 'n':
			return(4);
			break;
		default:
			return(4);
			break;
	}
}


//---------------------------------------------------------------------------
// * BaseIdx
//---------------------------------------------------------------------------
// 
int
HMM::BaseIdx2(char c)
{
	switch(c){
		case '1':
			return(0);
			break;
		case '2':
			return(1);
			break;
		case '3':
			return(2);
			break;
		case '4':
			return(3);
			break;
		case '5':
			return(4);
			break;
		case '6':
			return(5);
			break;
		default:
			cout<<"Invalid character: "<<c<<endl;
			return(0);
			break;
	}
}

//---------------------------------------------------------------------------
// * SetupForPrediction
//---------------------------------------------------------------------------
// An HMM consists of the following components: A set of states Q, an alphabet M, 
// a set of emission probabilities E(m,q) for each symbol m in M for each state q in Q,  
// and a set of transition probabilities T(q1,q2) for each pair of states q1 and q2 
// in Q.  The emission probabilities are stored in a table with one row for each 
// state q and one column for each symbol in the alphabet.  The transition probabilities
// are stored in a table with one row for each state q and one column for each state
// q.   An alphabet M of size |M| = n is represented by the integers 0-(n-1).  It 
// is the responsibility of the application to translate sequences of characters or
// other format data to and from this integer representation. 
//
void
HMM::SetupForPrediction(string &EmissionName,string &TransitionName,
		   Sequence &Sequence,vector<string> &StateNames)
{
  	mEmissionPr.Read(EmissionName);
	mTransitionPr.Read(TransitionName);
		
	mSequence = Sequence;
	mStateNames = StateNames;
	
	mNumSymbols = mEmissionPr.hsize();
	mNumStates = mTransitionPr.vsize();
	mSeqLength = mSequence.size();
}

//------------------------------------------------------
// * LogTransform
//------------------------------------------------------
// Does a log transform of the given table of values. 
//     
void 
HMM::LogTransform(Table<double> &DTable)
{
  int r,c;
  for(r = 0;r<DTable.vsize();r++){
  	for(c = 0;c<DTable.hsize();c++){
  		double raw = DTable[r][c];
  		DTable[r][c] = log10(raw);
  	}
  }
}


//---------------------------------------------------------------------------
// * LogTransform
//---------------------------------------------------------------------------
// Does a log transform of the emission and transition probabilities. 
// 
void
HMM::LogTransform()
{	
	// Convert probabilities into log(Pr).  All algorithms here
	// work in log space to minimize possibility of overflow errors. 
	LogTransform(mEmissionPr);
	LogTransform(mTransitionPr);	
}

//---------------------------------------------------------------------------
// * Viterbi
//---------------------------------------------------------------------------
// A dynamic programming algorithm to find most probable path through an HMM 
// given a set of observations.  The HMM is represented by a table of 
// transition probabilities and a table of emission probabilities.  It is 
// assumed that state 0 is both the beginning and the end state.  
//
// Note:  All probability calculations are performed in log space in order
// to minimize problems with under and overflow (i.e. probability products
// become sums in log space). 
//
void
HMM::Viterbi()
{
	mV.Allocate(mNumStates,mSeqLength+1);
	mPtr.Allocate(mNumStates,mSeqLength+1);
	// KJD Note: This initialiation is not really needed. 
	for(int r = 0;r<mV.mRows;r++){
		for(int c = 0;c<mV.mCols;c++){
			mV[r][c] = ROUGHLY_ZERO;
		}
	}
	mV[0][0] = 0; // log10(1)

	double MaxStatePr,StatePr;
	int MaxStateIdx;
	for(int oIdx = 1;oIdx<= mSeqLength;oIdx++){
		for(int sIdx = 0;sIdx < mNumStates;sIdx++){
			MaxStatePr = ROUGHLY_ZERO;// i.e. Pr~0
			for(int sIdx2 = 0;sIdx2<mNumStates;sIdx2++){
				// mV and mTransitionPr are in log space, so addition
				// is equivalent to multiplying probabilities
				StatePr = mV[sIdx2][oIdx-1]+mTransitionPr[sIdx2][sIdx];
				if (StatePr > MaxStatePr){
					MaxStatePr = StatePr;
					MaxStateIdx = sIdx2;
				}
			}
			// mEmissionPr is in log space, so addition is equivalent
			// to multiplying underlying probabilities.  
			mV[sIdx][oIdx]=
				mEmissionPr[sIdx][BaseIdx(mSequence[oIdx-1])]+MaxStatePr;
			mPtr[sIdx][oIdx]=MaxStateIdx;
		}
	}
	
	MaxStatePr = 0;
	for(int sIdx = 0;sIdx<mNumStates;sIdx++){
		// 0 = begin AND end.  log space again. 
		StatePr = mV[sIdx][mSeqLength]+mTransitionPr[sIdx][0]; 
		if (StatePr > MaxStatePr){
			MaxStatePr = StatePr;
			MaxStateIdx = sIdx;
		}
	}
	mPathProbability = MaxStatePr;
	mTracebackStart = MaxStateIdx;
}

//---------------------------------------------------------------------------
// * GetMostProbablePath
//---------------------------------------------------------------------------
// Performs a traceback through the state vector to recover the most 
// probable path of the HMM.  
//
void
HMM::GetMostProbablePath(vector<int> &path)
{
	int oIdx = mSeqLength;  // There is one more column than seq chars
	int sIdx = mNumStates-1;
	
	for(int i = 0;i<mSeqLength;i++){
		path.push_back(-1);
	}
	
	do{
		sIdx = mPtr[sIdx][oIdx];
		oIdx--;
		path[oIdx]=sIdx;
	}while(oIdx >1);
}


//---------------------------------------------------------------------------
// * UpdateCounts
//---------------------------------------------------------------------------
// Uses the sequence and state lables given to update the transition and 
// emission probabilities.  Values are added to the transition and emission 
// counts, so that this function can be called multiple times, updating 
// the tables as it goes.  
//
void 
HMM::UpdateCounts(Sequence Seq,vector<int> Labels)
{

  for(int i = 0;i< Seq.size();i++){
	 int state = Labels[i];
	 int letter = BaseIdx(Seq[i]);
	 mTransitionCount[state][letter]++;
	 mEmissionCount[state][letter]++;
  }
}


//---------------------------------------------------------------------------
// * UpdateEsitmateFromCounts
//---------------------------------------------------------------------------
//
void
HMM::UpdateEstimateFromCounts()
{
  double sumTransition = 0;
  double sumEmission = 0;

  // Transition probabilities. 
  for(int k = 0;k<mNumStates;k++){
	 // compute the denominator-> sum of transitions from k to each other state j 
	 for(int j = 0;j< mNumStates;j++){
		sumTransition += mTransitionCount[k][j];
	 }
	 // compute the probability for each k-i transition, by dividing count[k][i] by the 
	 // sum of transitions [k][j].  
	 for(int i = 0;i<mNumStates;i++){
		mTransitionPr[k][i] = mTransitionCount[k][i]/sumTransition;		
	 }
	 sumTransition = 0;
  }

  // Emission probabilities. 
  for(int k = 0;k<mNumStates;k++){
	 for(int b = 0;b< mNumSymbols;b++){
		sumEmission+= mEmissionCount[k][b];
	 }
	 for(int b = 0;b< mNumSymbols;b++){
		mEmissionPr[k][b] = mEmissionCount[k][b]/sumEmission;
	 }
	 sumEmission = 0;
  }

}


//---------------------------------------------------------------------------
// * GenerateSequenceFromModel
//---------------------------------------------------------------------------
//  Uses the HMM emmission and transition probabilities to generate a 
//  random sequence of the given length. 
//
void
HMM::GenerateSequenceFromModel(Sequence &SeqOut,vector<int> &StateOut,int size)
{
	RandomNumberGenerator RNG;
	RNG.SeedWithClock();
	
	for(int i = 0;i< size;i++){
		SeqOut.push_back('.');
		StateOut.push_back(0);
	}
	
	int currentState = 0,state;
	int seqIdx = 0;
	char emissionSymbol = '.';
	
	do{
		// Determine if there is a state transition. 
		double rndNum = RNG.Uniform(0,1);
		double rndLimit =0;
		for(state = 0;state<mNumStates;state++){
			rndLimit +=  mTransitionPr[currentState][state];
			if(rndNum < rndLimit){
				currentState = state;
				state = mNumStates;
			}
		}
		
		// Determine what symbol is emitted. 
		rndLimit = 0;
		rndNum = RNG.Uniform(0,1);
		for(int symbol = 0;symbol<mNumSymbols;symbol++){
			rndLimit +=  mEmissionPr[currentState][symbol];
			if(rndNum < rndLimit){
				emissionSymbol = symbol;
				symbol = mNumSymbols;
			}
		}
		
		SeqOut[seqIdx]=bugChar(emissionSymbol);
		StateOut[seqIdx] = currentState;
		seqIdx++;
		
	}while(seqIdx < size);
}










